import os
import os.path
import sqlite3
import ocrypt


class DBError (Exception):
    pass


acct_table_def = '''create table accounts (username text UNIQUE PRIMARY KEY,
                                           salt blob NOT NULL,
                                           verifier blob NOT NULL,
                                           is_locked integer,
                                           expires_at integer,
                                           is_admin   integer)'''

key_table_def = '''create table rsa_keys (keyname text UNIQUE PRIMARY KEY,
                                          pri_key blob NOT NULL,
                                          pub_key blob NOT NULL)'''

srp_table_def = '''create table srp_params (keyname text UNIQUE PRIMARY KEY,
                                            hash_type integer,
                                            prime_type integer)'''


class DB (object):
    
    def __init__(self, db_filename, allow_creation = False, hash_type = None, prime_type = None):
        do_create = db_filename == ':memory:' or not os.path.exists(db_filename)

        if do_create:
            if allow_creation:
                try:
                    fd = os.open(db_filename, os.O_WRONLY | os.O_EXCL | os.O_CREAT, 0600)
                    os.close(fd)
                except OSError, err:
                    raise Exception('Secure creation of database file failed: %s' % str(err))
            else:
                raise DBError('Credential database does not exist. Create it with `srp_admin create_credential_database %s`' % db_filename)
        
        self.con  = sqlite3.connect(db_filename)
        self.con.text_factory = str
        self.c    = self.con.cursor()

        if do_create:

            pri_str, pub_str = ocrypt.generate_keypair( 2048 )
            
            self.c.execute(acct_table_def)
            self.c.execute(key_table_def)
            self.c.execute(srp_table_def)

            self.c.execute('insert into rsa_keys values (?,?,?)', ('default', sqlite3.Binary(pri_str), sqlite3.Binary(pub_str)))
            self.c.execute('insert into srp_params values (?,?,?)',('default', hash_type, prime_type))
            self.con.commit()

        self.c.execute('select pri_key, pub_key from rsa_keys where keyname = ?', ('default',))
        
        tpl = self.c.fetchone()
        if not tpl:
            raise DBError('Unable to obtain RSA keys from database')
        
        self.pri_key_bytes = str(tpl[0])
        self.pub_key_bytes = str(tpl[1])

        self.c.execute('select hash_type, prime_type from srp_params where keyname = ?', ('default',))
        
        tpl = self.c.fetchone()
        if not tpl:
            raise DBError('Unable to obtain SRP Parameters from database')
        
        self.srp_hash_type  = tpl[0]
        self.srp_prime_type = tpl[1]
            

                        
    def close(self):
        self.con.close()

        
    def add_user(self, username, salt, verifier, is_locked=False, expires_at=-1, is_admin=False):
        locked = 1 if is_locked else 0
        admin  = 1 if is_admin else 0
        try:
            self.c.execute('insert into accounts values (?,?,?,?,?,?)', (username,sqlite3.Binary(salt),sqlite3.Binary(verifier), locked, expires_at, is_admin))
            self.con.commit()
        except sqlite3.IntegrityError:
            raise DBError('User already exists: %s' % username)
    
        
    def delete_user(self, username):
        try:
            self.c.execute('delete from accounts where username=?', (username,))
            self.con.commit()
        except sqlite3.Error, e:
            raise DBError('Failed to delete user: ' + str(e))
        
        
    def change_user_verifier(self, username, salt, verifier):
        self.c.execute('update accounts set salt=?, verifier=? where username=?', (sqlite3.Binary(salt),sqlite3.Binary(verifier),username))
        self.con.commit()
        
        
    def set_user_is_locked(self, username, b):
        self.c.execute('update accounts set is_locked=? where username=?', (1 if b else 0, username))
        self.con.commit()
        
        
    def set_user_expires_at(self, username, expires_at):
        self.c.execute('update accounts set expires_at=? where username=?', (expires_at, username))
        self.con.commit()
        
        
    def set_user_is_admin(self, username, b):
        self.c.execute('update accounts set is_admin=? where username=?', (1 if b else 0, username))
        self.con.commit()
        
        
    def get_user(self, username):
        self.c.execute('select salt, verifier, is_locked, expires_at, is_admin from accounts where username = ?', (username,))
        tpl = self.c.fetchone()
        if not tpl:
            return (None, None, None, None, None)
        else:
            return (str(tpl[0]), str(tpl[1]), bool(tpl[2]), tpl[3], bool(tpl[4]))
        
    
    def get_user_list(self):
        self.c.execute('select username from accounts')
        return [ str(u[0]) for u in self.c.fetchall() ]

    
    
    
if __name__ == '__main__':
    db = DB(':memory:')
    db.add_user( 'foo', 'bar', 'baz' )
    print 'get(foo): ', db.get_user('foo')
    db.update_user( 'foo', 'william', 'wallace')
    print 'get(foo): ', db.get_user('foo')
    print 'get(bin): ', db.get_user('bin')
    db.update_user( 'oops', 'william', 'wallace')
    db.add_user( 'foo', 'bar', 'baz' )
